perf: avoid full-vocab all-gather for draft greedy sampling - #1517

Motivation and scope

Draft/MTP proposal generation only consumes a greedy global winner and, where confidence is required, that winner's exact full-vocabulary probability. Gathering every vocabulary logit onto every TP rank wastes communication and memory. This PR changes draft models only; target-model sampling and its dense fallback are unchanged.

Final head 7a860c919731af5b24ee90bb7226a46a995ba4c9 is rebased onto the then-current upstream/main@3e0d9f70ce3e1070504ac812606aff7be6744b5a, which contains merged PRs #1509 and #1513.

Design and invariants

• Each TP rank computes local max, log-sum-exp, and argmax, then performs one fixed-size all-gather of three FP32 words per token.
• Global ties deterministically choose the smallest global token ID; IDs are bit-preserved in the FP32 communication buffer.
• ModelOutput.logits remains a logit tensor. Global IDs and normalization are explicit logits_token_ids and logits_logsumexp metadata.
• Exact selected-token probability is exp(selected_logit - full_vocab_logsumexp).
• CUDA Graph pad/unpad, Eagle accepted-row selection, dynamic MTP, and DSpark confidence handling preserve the metadata.
• Dense target-model sampling behavior is unchanged.

For Qwen3.5-27B (vocab_size=248320, TP4, BF16 dense logits), the all-gather output per token changes from 248320 × 2 = 496640 bytes to 4 × 3 × 4 = 48 bytes: 10,346.7× less collective payload (99.9903%). This follows directly from the actual tensor shapes and dtypes exercised by the CUDA tests and running service.

Final H100 correctness

• focused local suite: 39 passed
• focused H100 CUDA suite: 39 passed
• tests cover local/cross-rank ties, exact IDs/logits/log-sum-exp/log-probability against a dense full-vocabulary reference, CUDA Graph metadata pad/unpad, Eagle selected rows, DSpark output handling, and dense fallback semantics
• final service loaded source SHA256 d06cb14f015229a209b9af64700a0f3eca818a2667e83d06875fa87c5c9ecf04
• Black, Flake8, repository pre-commit, and git diff --check pass
• no reviews, comments, or required checks remain unresolved

H100 memory and end-to-end results

The matched TP4/MTP3 service A/B used the same four H100 80GB GPUs, Qwen3.5-27B files, FP8 weights, CUDA Graph settings/cache, source mount strategy, environment, and AIPerf streaming ISL=256/OSL=1024 corpus. Each formal candidate point is the median of three runs and all 42 runs completed with zero request errors.

Geometric-mean output-throughput gain is +8.55%. Relative to the immediate 3e0d9f70 control, #1517 adds about +1.5% geometric-mean throughput and +3.9% at C64; C2 is noisy and is not claimed as an isolated improvement.

Matched memory samples show:

• ready state: 67617 → 67523 MiB/GPU (94 MiB/GPU saved)
• C64 peak at 200 ms sampling: 69179 → 69105 MiB/GPU (74 MiB/GPU saved)

The end-to-end A/B uses the TP4 eagle_with_att MTP3 engine. DSpark-specific confidence semantics are established by the focused tests above; a full multi-node 3P1D DSpark deployment is downstream release validation and is not misrepresented as having been run by this PR.

Experiment records

Final candidate:

• launch: 260829-011244-2688792-sudo-n-docker-run-d-name-pr1517-final-7a860c91-p
• labeled warmup: 260829-011605-2699042-home-devsft-aiperf-venv-bin-aiperf-profile-model
• C1/C2/C4/C8/C16/C32: 260829-011645-2701809, 260829-011742-2703458, 260829-011823-2704888, 260829-011917-2706518, 260829-012000-2708554, 260829-012047-2711649
• C64 plus memory: 260829-012222-2717180-bash-s
• log archive/shutdown: 260829-012354-2723368-sh-c-sudo-n-docker-logs-pr1517-final-7a860c91-p1

Matched control memory:

• successful relaunch: 260829-012505-2724728-sudo-n-docker-run-d-name-pr1517-control-0be08549
• C64 memory: 260829-012806-2734716-bash-s
• log archive/shutdown: 260829-012921-2738290-sh-c-sudo-n-docker-logs-pr1517-control-0be08549-

Raw artifacts are retained under /home/devsft/qwen35_goal_20260829/pr1517_final_results and /home/devsft/qwen35_goal_20260829/pr1517_control_memory. Every performance command was wrapped by exp -m.